import io
import uuid
import cv2
from flask_socketio import emit
from common.yolo_model import model
from result.result import Result
from result.result_code import ResultCode
import os
# from ultralytics import YOLOv10
from utils.avi_to_mp4 import avi_to_mp4
import base64
from PIL import Image
import json
# from flask import get_json()
from model.detect_result import DetectResult
from flask_jwt_extended import get_jwt_identity
from datetime import datetime
from common.database import db


class YoloService(object):

    @staticmethod
    def yolo(data):
        results = None
        save_dir = '.\\static\\detect'
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)
        name = str(uuid.uuid4())
        try:
            if data['type'] == 'image':
                results = model.predict(source=os.path.join('.\\static\\upload', data['image']),
                                        imgsz=data['image_size'],
                                        conf=data['conf_threshold'], iou=data['iou_threshold'], max_det=data['max_det'],
                                        augment=data['augment'], agnostic_nms=data['agnostic_nms'],
                                        retina_masks=data['retina_masks'], classes=data['classes'],
                                        save=True, project=save_dir, name=name)
                # print(results)
            elif data['type'] == 'video':
                results = model.predict(source=os.path.join('.\\static\\upload', data['video']),
                                        imgsz=data['image_size'],
                                        conf=data['conf_threshold'], iou=data['iou_threshold'], batch=data['batch'],
                                        max_det=data['max_det'], vid_stride=data['vid_stride'],
                                        augment=data['augment'], agnostic_nms=data['agnostic_nms'],
                                        retina_masks=data['retina_masks'], classes=data['classes'],
                                        save=True, project=save_dir, name=name)
                if not avi_to_mp4(os.path.join(save_dir, name)):
                    os.remove(os.path.join(save_dir, name))
                    return Result.error(ResultCode.INTERNAL_SERVER_ERROR.value, "服务器在将avi格式转成mp4格式时出现了一点错误")
        except Exception as e:
            print(e)
            return Result.error(ResultCode.INTERNAL_SERVER_ERROR.value, "检测出错")
        new_results = {
            'result': [],
            'save_dir': results[0].save_dir.split('\\')[-1],
            'speed': {'inference': 0.0,
                      'postprocess': 0.0,
                      'preprocess': 0.0
                      }
        }

        for result in results:
            new_results['result'].append(json.loads(result.tojson()))
            new_results['speed']['inference'] += result.speed['inference']
            new_results['speed']['postprocess'] += result.speed['postprocess']
            new_results['speed']['preprocess'] += result.speed['preprocess']

        # new_results['result'] = new_results['result'][0]
        # 新增结果到数据库
# Python 中没有三元运算符 ? :，需要使用 if-else 语句来替代
        upload_path = data['image'] if len(
            data['image']) != 0 else data['video']
        speed = results[0].speed['inference'] + \
            results[0].speed['postprocess'] + results[0].speed['preprocess']
        bath = data['batch'] if data['type'] == 'video' else None
        vid_stride = data['vid_stride'] if data['type'] == 'video' else None
        try:
            detect_result = DetectResult(uuid=uuid.uuid4(), user_email=get_jwt_identity(), detect_time=datetime.now(),
                                         upload_path=upload_path, result_path=new_results['save_dir'],
                                         speed=speed, imgsz=data['image_size'], conf=data['conf_threshold'],
                                         iou=data['iou_threshold'], max_det=data['max_det'], augment=data['augment'],
                                         agnostic_nms=data['agnostic_nms'], retina_masks=data['retina_masks'],
                                         classes=json.dumps(data['classes']),
                                         batch=bath, vid_stride=vid_stride)

            db.session.add(detect_result)
            db.session.commit()
        except Exception as e:
            print(e)
            return Result.error(ResultCode.DATABASE_ERROR.value, "数据库写入错误")
        return Result.success_data(new_results)

    @staticmethod
    def handle_video_frame(frame_data, sid):
        # print(frame_data['flag'])
        if not frame_data['flag']:
            emit('processed-image', {
                'image': frame_data['data']
                # "image": f"data:image/webp;base64,{frame_data['data']}"
            }, room=sid)
            # print("发送了数据")
            # emit('processed-image', "hello world", room=sid)
            return
        try:
            # frame_bytes = base64.b64decode(frame_data)
            image = Image.open(io.BytesIO(frame_data['data']))
        except base64.binascii.Error as e:
            print(f"Base64 解码出错: {e}")
            return
        except Image.UnidentifiedImageError as e:
            print(f"无法识别图像数据: {e}")
            return e
        except Exception as e:
            print(f"其他错误: {e}")
            return e

        # model = YOLOv10(model_path)
        try:
            results = model.predict(
                image, imgsz=320, conf=0.25, batch=8, save=False)
        except Exception as e:
            print(f"图像检测出错：{e}")
            return e
        finally:
            image.close()

        try:
            annotated_image = results[0].plot()
            retval, buffer = cv2.imencode('.webp', annotated_image)
            image_bytes = buffer.tobytes()

            base64_image = base64.b64encode(image_bytes).decode('utf-8')
            response_data = {
                'image': f'data:image/webp;base64,{base64_image}'
            }
        except Exception as e:
            print(e)
            return e
        emit('processed-image', response_data, room=sid)
        # emit('processed-image', frame_data)
